#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import print_function
import rospy
import csv
import math
from geometry_msgs.msg import Point
from std_msgs.msg import Float32MultiArray
import os
import sys

# add path: /home/chiweiming/catkin_ws/src/cwm_ros
sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))

#---Directory Path---#
# Path: /home/chiweiming/catkin_ws/src/cwm_ros/script
dirPath = os.path.dirname(os.path.realpath(__file__))

# 设置输出的名字
savename = 'yzg-brain_auto_A'

# Number of experiment
numexp = '1'

# 定义全局变量来存储两个目标点
tb3_1_target = None
tb3_2_target = None

# 打开 CSV 文件进行写入
f = open(dirPath + '/ExperimentalData/' + savename + '-' + numexp + '.csv', 'w')
csv_writer = csv.writer(f)
csv_writer.writerow([
    "Leader X", "Leader Y", "Leader V_x", "Leader W", 
    "follower1 X", "follower1 Y", "follower1 V_x", "follower1 W", 
    "follower2 X", "follower2 Y", "follower2 V_x", "follower2 W", 
    "state", "tb3_1_target_x", "tb3_1_target_y", "tb3_2_target_x", "tb3_2_target_y"
])

def myhook():
    f.close()
    print("shutdown time!")

# 领航者位置回调函数
def leaderOdom_CB(data):
    # 控制接收频率 
    rate = rospy.Rate(10)
    rate.sleep()

    if (abs(data.data[2]) <= 0.3) and (abs(data.data[3]) <= 1.0):
        # 创建列表用于存储所需数据: 领航者的位置（x，y） 速度（v，w）
        a = data.data[:18]  # 领航者的数据
        # 如果两个目标点都已经接收到，就记录到 CSV
        if tb3_1_target is not None and tb3_2_target is not None:
            recordTargetsToCSV(a)

# 记录目标点信息到 CSV 文件
def recordTargetsToCSV(leader_data):
    global tb3_1_target, tb3_2_target
    # 将 leader_data 转换为 list 并合并目标坐标
    leader_data_list = list(leader_data)
    # 写入数据到 CSV 文件
    csv_writer.writerow(leader_data_list + [tb3_1_target.x, tb3_1_target.y, tb3_2_target.x, tb3_2_target.y])
    rospy.loginfo("Recorded: Leader data and Targets (%f, %f), (%f, %f)", 
                  tb3_1_target.x, tb3_1_target.y, tb3_2_target.x, tb3_2_target.y)
    # 清空目标点信息，防止重复记录
    tb3_1_target = None
    tb3_2_target = None


# 订阅 tb3_1 目标点信息
def formationTargetCallback1(msg):
    global tb3_1_target
    tb3_1_target = msg
    rospy.loginfo("Received tb3_1 target: x = %f, y = %f", msg.x, msg.y)

# 订阅 tb3_2 目标点信息
def formationTargetCallback2(msg):
    global tb3_2_target
    tb3_2_target = msg
    rospy.loginfo("Received tb3_2 target: x = %f, y = %f", msg.x, msg.y)

# 监听数据
def listenTOLeader():
    # 初始化节点
    rospy.init_node('listenTOLeader')
    
    # 订阅领航者的位置话题
    rospy.Subscriber('/publishAllLocations', Float32MultiArray, leaderOdom_CB, queue_size=1, buff_size=52428800)
    
    # 订阅两个跟随者的目标点话题
    rospy.Subscriber("tb3_1/formation_target", Point, formationTargetCallback1)
    rospy.Subscriber("tb3_2/formation_target", Point, formationTargetCallback2)
    
    rospy.spin()  # 保持节点运行

if __name__ == '__main__':
    try:
        listenTOLeader()
    except rospy.ROSInterruptException:
        pass
    finally:
        # 关闭 CSV 文件
        f.close()
